package org.charlott2.fruitrec

import android.graphics.Bitmap
import android.util.Log
import androidx.activity.ComponentActivity
import org.json.JSONArray
import org.json.JSONObject
import org.pytorch.IValue
import org.pytorch.LiteModuleLoader
import org.pytorch.Module
import org.pytorch.torchvision.TensorImageUtils
import java.io.File
import java.io.FileOutputStream
import java.nio.charset.Charset

class AIManager private constructor() {

    /**
     * 单例模式
     */
    companion object {
        private var instance: AIManager? = null
        fun getInstance(): AIManager {
            return instance ?: AIManager()
        }
    }

    // 定义模型
    private var model: Module? = null

    // 定义标签列表，下标是它对应的值
    private var labels: Array<String> = emptyArray()

    // 定义价格映射表
    private var prices: Map<String, String> = emptyMap()

    init {
        instance = this
    }

    fun init(ctx: ComponentActivity) {
        // 读取模型
        val modelPath = assetFilePath(ctx, "model.ptl")
        model = LiteModuleLoader.load(modelPath)
        Log.d("aa", "model loaded successfully from $modelPath")

        // 读取标签
        val labelsPath = assetFilePath(ctx, "labels.json")
        labels = loadLabelsFromJson(labelsPath)
        Log.d("aa", "labels loaded successfully from $labelsPath, len: ${labels.size}")

        // 读取价格
        val pricesPath = assetFilePath(ctx, "prices.json")
        prices = loadPricesFromJson(pricesPath)
        Log.d("aa", "prices loaded successfully from $pricesPath, len: ${prices.size}")
    }

    // 将文件从 assets 文件夹中读取到本地，返回文件绝对路径
    private fun assetFilePath(ctx: ComponentActivity, assetName: String): String {
        Log.d("aa", ctx.filesDir.toString())
        // 根据文件名，构建完整文件对象，返回文件绝对路径
        val file = File(ctx.filesDir, assetName)
        if (!file.exists()) {
            // 从 assets 文件夹中读取文件
            ctx.assets.open(assetName).use { inputStream ->
                FileOutputStream(file).use { outputStream ->
                    inputStream.copyTo(outputStream)
                }
            }
        }
        return file.absolutePath
    }

    private fun loadLabelsFromJson(filePath: String): Array<String> {
        val jsonStr = File(filePath).readText(Charset.defaultCharset())
        val jsonArray = JSONArray(jsonStr)
        return Array(jsonArray.length()) { jsonArray.getString(it) }
    }

    private fun loadPricesFromJson(filePath: String): Map<String, String> {
        val jsonStr = File(filePath).readText(Charset.defaultCharset())
        val jsonObject = JSONObject(jsonStr)
        val pricesMap = mutableMapOf<String, String>()
        jsonObject.keys().forEach { key ->
            pricesMap[key] = jsonObject.getString(key)
        }
        return pricesMap
    }

    // 预测
    fun inference(bitmap: Bitmap): Map<String, String> {
        // 将图片转换为100*100
        val resizeBm = Bitmap.createScaledBitmap(bitmap, 100, 100, false)
        // 将图片转换为张量
        val imageTensor = TensorImageUtils.bitmapToFloat32Tensor(
            resizeBm,
            floatArrayOf(0.0f, 0.0f, 0.0f), // 表示不进行均值处理
            floatArrayOf(1.0f, 1.0f, 1.0f) // 表示不进行标准差处理
        )
        // 模型预测
        val outputTensor = model!!.forward(IValue.from(imageTensor)).toTensor()
        // 将预测结果转换为数组
        val scores = outputTensor.dataAsFloatArray
        Log.d("aa", "scores: ${scores.contentToString()}")

        // 找出最大值，得到预测结果
        val idx = scores.indices.maxByOrNull { scores[it] } ?: -1
        val type = if (idx == -1) "未知" else labels[idx]
        val price = prices[type] ?: "0.0"
        return mapOf("type" to type, "price" to price)
    }

}
